from flask import Flask, request, jsonify, Response, send_from_directory
from datetime import datetime
import os
import threading
from run_uvr import *
from ModelConfig import GenerationConfig
from song_constant import *
from infer_combine import run_task_infer, call_generation
from run_trainer import *
import glob

app = Flask(__name__)

# 设置上传文件的目录
UPLOAD_FOLDER = 'uploads/'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)


# 确保文件名是安全的
def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in {'mp3'}


@app.route('/upload_mp3/<task_id>', methods=['POST'])
def upload_mp3(task_id):
    # 检查是否有文件被上传
    if 'file' not in request.files:
        return jsonify({'message': '没有训练数据文件'}), 400
    file = request.files['file']
    # 如果用户没有选择文件，浏览器也会提交一个无文件的文件对象
    if file.filename == '':
        return jsonify({'message': 'No selected file'}), 400
    if file and allowed_file(file.filename):
        # 保存文件到指定目录
        # filename = secure_filename(file.filename)
        filename = file.filename
        print('filename', filename, file.filename)
        file.save(os.path.join(UPLOAD_FOLDER, task_id + filename))
        return jsonify({'message': f'File uploaded successfully', 'filename': filename}), 200
    else:
        return jsonify({'message': 'File type not allowed'}), 400


def infer_model():
    pass


@app.route('/check_status/<task_id>', methods=['GET'])
def check_status(task_id):
    # 检查完整的.pth文件是否存在
    full_path_file = f'/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/weights/{task_id}.pth'
    training_files = glob.glob(f'/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/weights/{task_id}*.pth')
    if os.path.exists(full_path_file):
        return jsonify({'status': 'trained', 'info': training_files})
    # 检查是否有正在训练的.pth文件
    training_logs_path =  f'/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/logs/{task_id}'
    if os.path.exists(training_logs_path):
        return jsonify({'status': 'training', 'info': '训练任务已经启动。'})
    if training_files:
        return jsonify({'status': 'training', 'info': training_files})
    # 如果两者都不存在，返回错误状态
    return jsonify({'status': 'error','info':'预处理中'})


# 定义一个函数来处理模型训练
def train_model(task_id, vocal, file_path, folder_name, format0):
    # 这里是模型训练的伪代码
    print(f'开始训练 {task_id}  【{folder_name}】')

    # 训练模型...
    # 假设训练成功
    instru = task_id + '_instru'
    os.makedirs(instru, exist_ok=True)
    run_model(os.path.abspath(folder_name), os.path.abspath(vocal), os.path.abspath(instru), format0)
    print('训练数据UVR后 在', instru)
    perform_training_task(task_id,instru,'rmvpe_gpu')

    return 'Model trained successfully'


@app.route('/')
def hello():
    html_content = """
    <html>
        <head>
            <title>Hello Page</title>
        </head>
        <body>
            <h1>Hello, World!</h1>
        </body>
    </html>
    """
    return html_content


@app.route('/train/<task_name>', methods=['POST'])
def train(task_name):
    # 获取上传的文件
    print('in train')

    if 'file' not in request.files:
        print('no datas')
        return jsonify({'message': '没有训练数据文件'}), 400
    file = request.files['file']

    print('receive', file.filename, file.name)
    # 生成文件夹名，格式为userid-datetime
    userid = 'user123'  # 假设的userid
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    folder_name = f'training_{userid}_{timestamp}'
    task_id = f'training_{userid}_{timestamp}'
    folder_path = folder_name  # os.path.join('/path/to/save/models', folder_name)  # 指定保存路径

    # 保存文件到对应的文件夹
    os.makedirs(folder_path, exist_ok=True)
    print('文件保存到', folder_path)
    path_ = os.path.join(folder_path, file.filename)
    file.save(path_)
    ext = os.path.splitext(file.name)[-1]
    vocal = task_id + '_vocal'
    os.makedirs(vocal, exist_ok=True)
    # 创建一个线程来执行训练任务
    train_thread = threading.Thread(target=train_model, args=(task_id, vocal, folder_path, folder_name, ext))
    train_thread.start()

    # 立即返回状态和文件夹名，不等待训练完成
    return jsonify({
        'status': 'training',
        'folder_name': folder_name,
        'task_id': task_id  # 使用线程的名称作为任务ID
    })


@app.route('/run_rvc', methods=['POST'])
def run_rvc():
    # 必须要taskid ，taskid就是model id
    print('run rvc')
    data = request.json
    task_id = data.get('task_id')
    if task_id is None:
        return jsonify({'error': 'No taskid provided'})
    user_id = data.get('user_id', '__defualt_user_id')
    song_id = data.get('song_id', '__default_song_id')
    config = GenerationConfig('VC', user_id, GenerationConfig.get_ts())
    print('gen config')
    config.make_folder()
    print('start thread')
    train_thread = threading.Thread(target=call_generation, args=(task_id, user_id, song_id, config))
    train_thread.start()
    # config = call_generation(task_id, user_id, song_id)
    print('return status')

    return jsonify({
        'status': 'vc_runing',
        'ai_vocal': config.vocal_ai_gen_path,
        'song': config.song_gen_path,
        'task_id': task_id  # 使用线程的名称作为任务ID
    })
    # return Response(open(config.song_gen_path, 'rb'), mimetype='application/octet-stream')


@app.route('/get_prediction/<output_dir>', methods=['GET'])
def get_prediction(output_dir):
    print('<预测>', output_dir)
    prediction_file_path = os.path.join(output_dir, 'output.wav')
    print('文件path', prediction_file_path)
    if os.path.exists(prediction_file_path):
        print('找到这个文件', prediction_file_path)
        return send_from_directory(os.path.dirname(prediction_file_path), os.path.basename(prediction_file_path),
                                   as_attachment=True)
    else:
        print('找不到这个文件', prediction_file_path)
        return jsonify({"error": "Prediction result not found"}), 404


@app.route('/infer', methods=['POST'])
def infer():
    # 获取上传的文件

    if 'file' not in request.files:
        return jsonify({'error': 'No file part in the request'})

    file = request.files['file']
    # model id
    taskid = request.args.get('taskid')
    if taskid is None:
        return jsonify({'error': 'No taskid provided'})

    if file.filename == '':
        return jsonify({'error': 'No selected file'})

    print('receive', file.filename, file.name)
    # 生成文件夹名，格式为userid-datetime
    userid = 'user123'  # 假设的userid
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    folder_name = f'{userid}_{timestamp}'
    task_id = f'{userid}_{timestamp}'
    folder_path = folder_name  # os.path.join('/path/to/save/models', folder_name)  # 指定保存路径

    # 保存文件到对应的文件夹
    os.makedirs(folder_path, exist_ok=True)
    file.save(os.path.join(folder_path, file.filename))
    ext = os.path.splitext(file.name)[-1]
    vocal = task_id + '_vocal'
    os.mkdir(vocal)
    # 创建一个线程来执行训练任务
    train_thread = threading.Thread(target=infer_model, args=(task_id, vocal, folder_path, folder_name, ext))
    train_thread.start()

    # 立即返回状态和文件夹名，不等待训练完成
    return jsonify({
        'status': 'training',
        'folder_name': folder_name,
        'task_id': task_id  # 使用线程的名称作为任务ID
    })


# 启动 Flask 应用
if __name__ == '__main__':
    app.run(debug=True, port=6006)